# 1. 导入需要的库
from docx import Document  # 读取Word文档（员工手册）
import chromadb  # 向量数据库
from chromadb.utils import embedding_functions  # 向量生成工具
from openai import OpenAI  # 调用OpenAI接口
import os

# --------------------------
# 配置信息（学生需要改这里！）
# --------------------------
OPENAI_API_KEY = "你的OpenAI API密钥"  # 替换成自己的密钥
WORD_FILE_PATH = "employee_manual.docx"  # 员工手册文件路径（确保和代码同文件夹）
CHROMA_COLLECTION_NAME = "employee_manual"  # 向量数据库的"集合名"（类似文件夹）


# 2. 初始化工具（向量生成、LLM调用、向量数据库）
# 2.1 初始化DashScope向量生成工具（把文本转成向量）
class DashScopeEmbeddingFunction:
    def __init__(self, batch_size=10):
        self.client = OpenAI(
            api_key=os.getenv("DASHSCOPE_API_KEY"),  # 如果您没有配置环境变量，请在此处用您的API Key进行替换
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"  # 百炼服务的base_url
        )
        self.batch_size = batch_size

    def __call__(self, input):
        if isinstance(input, str):
            input = [input]

        # 分批处理输入，避免超过批处理大小限制
        embeddings = []
        for i in range(0, len(input), self.batch_size):
            batch = input[i:i + self.batch_size]
            completion = self.client.embeddings.create(
                model="text-embedding-v4",
                input=batch,
                dimensions=1024,  # 指定向量维度（仅 text-embedding-v3及 text-embedding-v4支持该参数）
                encoding_format="float"
            )

            # 提取嵌入向量
            batch_embeddings = [item.embedding for item in completion.data]
            embeddings.extend(batch_embeddings)

        return embeddings


# 实例化DashScope嵌入函数
dashscope_ef = DashScopeEmbeddingFunction()

# 2.2 初始化OpenAI LLM（用来生成回答）
client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# 2.3 初始化Chroma向量数据库（本地存储，不用联网）
chroma_client = chromadb.Client()
# 先删除旧的集合（避免重复数据，方便重复运行）
try:
    chroma_client.delete_collection(name=CHROMA_COLLECTION_NAME)
except:
    pass
# 创建新集合（用来存员工手册的向量）
collection = chroma_client.create_collection(
    name=CHROMA_COLLECTION_NAME,
    embedding_function=dashscope_ef  # 绑定向量生成工具
)


# --------------------------
# 3. 第一步：处理员工手册（数据准备+存储）
# --------------------------
def load_employee_manual(file_path):
    """读取员工手册Word文件，提取文本内容（简易版：只读段落）"""
    doc = Document(file_path)
    all_text = []
    for para in doc.paragraphs:
        # 过滤空段落，保留有效文本
        clean_text = para.text.strip()
        if clean_text:
            all_text.append(clean_text)
    # TODO 1：文档解析优化——后续可处理手册中的表格（用doc.tables遍历表格）
    # TODO 2：文本清洗优化——去除特殊字符（如%、¡），示例：re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s。，]', '', clean_text)
    return "\n".join(all_text)  # 把所有段落拼成一个长文本


def split_text_simple(text, chunk_size=200):
    """简易文本切分：按固定长度切分（避免片段太长/太短）"""
    chunks = []
    for i in range(0, len(text), chunk_size):
        chunk = text[i:i + chunk_size]
        chunks.append(chunk)
    # TODO 3：切分策略优化——后续用语义切分（如LangChain的RecursiveCharacterTextSplitter），按句子断句
    return chunks


# 执行：读取→切分→存入向量库
manual_text = load_employee_manual(WORD_FILE_PATH)
text_chunks = split_text_simple(manual_text, chunk_size=200)

# 给每个片段加唯一ID（方便后续追溯）
chunk_ids = [f"chunk_{i}" for i in range(len(text_chunks))]

# 存入Chroma向量库（自动生成向量）
collection.add(
    documents=text_chunks,  # 文本片段
    ids=chunk_ids  # 唯一ID
    # TODO 4：元数据优化——后续可加元数据（如章节名、页码），示例：metadatas=[{"chapter":"薪酬福利", "page":11}]
)
print(f"员工手册处理完成！共生成 {len(text_chunks)} 个文本片段，已存入向量库")


# --------------------------
# 4. 第二步：用户问答（检索+生成）
# --------------------------
def rag_answer_question(question, top_k=3):
    """
    RAG问答核心函数：
    1. 问题向量化 2. 检索相似片段 3. 结合片段生成答案
    """
    # 1. 检索：找向量库中最相似的top_k个片段
    results = collection.query(
        query_texts=[question],  # 用户问题
        n_results=top_k  # 取前3个最相似片段
        # TODO 5：检索优化——后续可加元数据过滤（如只查“年假”相关章节）
    )
    # 提取检索到的文本片段（去重）
    retrieved_chunks = list(set(results["documents"][0]))
    # 构建上下文（把片段拼起来，告诉LLM“基于这些内容回答”）
    context = "\n".join([f"- {chunk}" for chunk in retrieved_chunks])
    print(f"\n【检索到的参考内容】\n{context}\n")

    # 2. 生成：调用LLM，结合上下文回答问题
    prompt = f"""
    你是公司HR助理，只能基于以下员工手册内容回答问题，不要编造信息：
    {context}

    用户问题：{question}
    要求：回答简洁明了，分点说明（如果有多个要点）。
    """
    # TODO 6：Prompt优化——后续可加角色约束、格式约束（如“用表格展示”）
    # TODO 7：LLM优化——后续可替换为开源LLM（如用ollama调用llama3），不用依赖OpenAI

    # 调用OpenAI生成回答
    response = client.chat.completions.create(
        model="qwen-plus",  # 使用阿里云百炼支持的qwen-plus模型
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2  # 温度低一点，回答更准确（少幻觉）
    )
    return response.choices[0].message.content


# --------------------------
# 5. 运行演示
# --------------------------
if __name__ == "__main__":
    # 示例问题（学生可替换成其他问题，如“技术序列T3职级薪资多少？”“加班有补偿吗？”）
    user_question = "员工工作满3年，年假有几天？需要提前多久申请？"
    print(f"【用户问题】{user_question}")

    # 生成答案
    answer = rag_answer_question(user_question)
    print(f"【RAG生成答案】\n{answer}")
